"""
FastAPI部署接口
功能：接收输入数据，返回损伤预测结果
"""
import sys
import os
import joblib
import numpy as np
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from models.multimodel_net import MultiModalNet

# 获取当前文件所在目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 获取项目根目录
project_root = os.path.dirname(current_dir)
# 将项目根目录添加到 sys.path
sys.path.insert(0, project_root)

app = FastAPI()
# 指定设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MultiModalNet().to(device)

# 检查模型文件是否存在
model_path = 'outputs/best_model.pth'
if not os.path.exists(model_path):
    raise FileNotFoundError(f"模型文件 '{model_path}' 未找到，请检查文件路径。")

# 加载模型
try:
    model.load_state_dict(torch.load(model_path, map_location=device))
except Exception as e:
    raise HTTPException(status_code=500, detail=f"加载模型时出现错误: {str(e)}")

model.eval()

# 检查 scaler 和 encoder 文件是否存在
scaler_path = 'outputs/scaler.pkl'
encoder_path = 'outputs/encoder.pkl'
if not os.path.exists(scaler_path) or not os.path.exists(encoder_path):
    raise FileNotFoundError(f"scaler 或 encoder 文件未找到，请检查文件路径。")

# 加载 scaler 和 encoder
try:
    scaler = joblib.load(scaler_path)
    encoder = joblib.load(encoder_path)
except Exception as e:
    raise HTTPException(status_code=500, detail=f"加载 scaler 或 encoder 时出现错误: {str(e)}")


class RequestData(BaseModel):
    acc_x: list  # X轴加速度（1000点）
    acc_y: list  # Y轴
    acc_z: list  # Z轴
    beam_width: float  # 梁宽（mm）
    beam_height: float  # 梁高（mm）
    rebar_ratio: float  # 配筋率
    concrete_grade: str  # 混凝土等级


@app.post("/predict")
async def predict(data: RequestData):
    try:
        # 预处理输入数据
        wave = np.stack([data.acc_x, data.acc_y, data.acc_z], axis=0)
        wave = (wave / np.max(np.abs(wave))).astype(np.float32)

        # 结构参数编码
        num_features = np.array([data.beam_width, data.beam_height, data.rebar_ratio]).reshape(1, -1)
        num_features_scaled = scaler.transform(num_features)
        encoded_grade = encoder.transform([[data.concrete_grade]])
        struct_params = np.hstack([num_features_scaled, encoded_grade]).astype(np.float32)

        # 转换为Tensor并移动到设备上
        wave_tensor = torch.tensor(wave[np.newaxis, ...]).to(device)
        struct_tensor = torch.tensor(struct_params).to(device)  # 确保维度正确

        print("Wave tensor shape:", wave_tensor.shape)
        print("Struct tensor shape:", struct_tensor.shape)

        # 进行预测
        with torch.no_grad():
            pred = model(wave_tensor, struct_tensor)

        print("Prediction:", pred.item())

        return {"damage_index": pred.item()}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"预测过程中出现错误: {str(e)}")